#!/usr/bin/env python
#-*- coding:utf-8 -*-

from ast import AST
import random
import math
from planner.planner import Planner
from graph.node_tree import NodeTree
from common.node import Node
from common.edge import Edge
from planner.astar_planner import AStarPlanner

class BiRRTPlanner(Planner):
    def __init__(self, graph=None, start=None, goal=None):
        super(BiRRTPlanner, self).__init__(graph, start, goal)
        self.connected_node = None
        self.heading_vec_dict = {}
    
    @property
    def connected_node(self):
        if not hasattr(self, "_connected_node"):
            self._connected_node = None
        return self._connected_node

    @connected_node.setter
    def connected_node(self, connected_node):
        self._connected_node = connected_node

    @property
    def start_tree(self):
        if not hasattr(self, "_start_tree"):
            self._start_tree = NodeTree()
        return self._start_tree

    @start_tree.setter
    def start_tree(self, start_tree):
        self._start_tree = start_tree

    @property
    def goal_tree(self):
        if not hasattr(self, "_goal_tree"):
            self._goal_tree = NodeTree()
        return self._goal_tree

    @goal_tree.setter
    def goal_tree(self, goal_tree):
        self._goal_tree = goal_tree

    @property
    def random_trees(self):
        return self.start_tree, self.goal_tree

    def block_free(self, from_node, to_node):
        return not self.graph.blocked(from_node, to_node)

    def almost_connected(self, first_node, second_node, threshold=10):
        if self.cost(first_node, second_node) < threshold:
            return True
        return False

    def _expand_node_and_edge(self, graph, from_node, to_node):
        graph.add_node(to_node)
        _weight = self.cost(from_node, to_node)
        graph.add_edge(Edge(from_node.nid, to_node.nid, _weight))
        _heading_vec = (to_node.pos[0] - from_node.pos[0], to_node.pos[1] - from_node.pos[1])
        self.heading_vec_dict[to_node.nid] = _heading_vec

    def _take_small_step(self, from_node, to_node, step_length=10):
        if from_node is None or to_node is None:
            return None
        dist = ((to_node.pos[0] - from_node.pos[0]) ** 2 + (to_node.pos[1] - from_node.pos[1]) ** 2) ** 0.5
        if dist == 0:
            return None
        if dist < step_length:
            return to_node
        rate = step_length / dist
        new_w = int(rate * (to_node.pos[0] - from_node.pos[0]) + from_node.pos[0])
        new_h = int(rate * (to_node.pos[1] - from_node.pos[1]) + from_node.pos[1])
        new_node = Node((new_w, new_h), (new_w, new_h))
        return new_node

    def theta(self, vec1, vec2):
        if vec1 == (0, 0):
            return math.atan2(vec2[1], vec2[0])
        if vec2 == (0, 0):
            return math.atan2(vec1[1], vec1[0])
        v1_mod = (vec1[0] * vec1[0] +  vec1[1] * vec1[1]) ** 0.5
        v2_mod = (vec2[0] * vec2[0] +  vec2[1] * vec2[1]) ** 0.5
        v1_v2_in_prod = vec1[0] * vec2[0] + vec1[1] * vec2[1]
        return math.acos(v1_v2_in_prod / (v1_mod * v2_mod + 0.000001))

    def motion_cost(self, from_node, to_node, dist_rate=0.4):
        _target_vec = (to_node.pos[0] - from_node.pos[0], to_node.pos[1] - from_node.pos[1])
        delta_dist = ((to_node.pos[0] - from_node.pos[0]) ** 2 + (to_node.pos[1] - from_node.pos[1]) ** 2 ) ** 0.5
        delta_theta = self.theta(self.heading_vec_dict[from_node.nid], _target_vec)
        cost_of_dist = max(1/(delta_dist + 0.0000001), delta_dist)
        cost_of_theta = abs(math.tan(delta_theta))
        total_cost = cost_of_dist * dist_rate + cost_of_theta * (1 - dist_rate)
        return total_cost

    def _find_nearest_node(self, tree, target, mode='dist'):
        nearst_dist = float('inf')
        # nearst_node = tree.root_node
        nearst_node = None
        # print("91", target)
        for node in tree.nodes():
            # print("92", node, tree.nodes())
            if mode == 'cost':
                cost_val = self.cost(node, target)
            elif mode == 'motion':
                cost_val = self.motion_cost(node, target)
            else: #dist
                cost_val = self.distance(node, target)
            # print("100", cost_val)
            if cost_val < nearst_dist:
                nearst_dist = cost_val
                nearst_node = node
        return nearst_node, nearst_dist
            
    def _random_sample(self):
        _w = random.randint(0, self.graph.width - 1)
        _h = random.randint(0, self.graph.height - 1)
        return Node((_w, _h), (_w, _h))

    def _init_process(self):
        self._tree_1st = NodeTree()
        self._tree_2nd = NodeTree()
        self.start_node = Node(self.start, self.start)
        self.goal_node = Node(self.goal, self.goal)
        self._tree_1st.add_root(self.start_node)
        self._tree_2nd.add_root(self.goal_node)
        start_vec = (self.goal_node.pos[0] - self.start_node.pos[0], self.goal_node.pos[1] - self.start_node.pos[1])
        self.heading_vec_dict[self.start] = start_vec
        goal_vec = (self.start_node.pos[0] - self.goal_node.pos[0], self.start_node.pos[1] - self.goal_node.pos[1])
        self.heading_vec_dict[self.goal] = goal_vec

    def _build_bi_tree(self, step_length=10, max_epoch=5000):
        self._init_process()
        for e in range(max_epoch):
            print("epoch {}".format(e))
            target_1st = self._random_sample()
            nearest_1st, near_cost = self._find_nearest_node(self._tree_1st, target_1st, mode='dist')
            if not nearest_1st or near_cost < 1:
                continue
            small_step_1st = self._take_small_step(nearest_1st, target_1st, step_length)
            if small_step_1st and self.block_free(nearest_1st, small_step_1st):
                self._expand_node_and_edge(self._tree_1st, nearest_1st, small_step_1st)
                nearest_2nd, _ = self._find_nearest_node(self._tree_2nd, small_step_1st, mode='dist')
                small_step_2nd = self._take_small_step(nearest_2nd, small_step_1st, step_length)
                if small_step_2nd and self.block_free(nearest_2nd, small_step_2nd):
                    self._expand_node_and_edge(self._tree_2nd, nearest_2nd, small_step_2nd)
                    while not self.almost_connected(small_step_2nd, small_step_1st):
                        small_step_tmp = self._take_small_step(small_step_2nd, small_step_1st, step_length)
                        if small_step_tmp and self.block_free(small_step_2nd, small_step_tmp):
                            self._expand_node_and_edge(self._tree_2nd, small_step_2nd, small_step_tmp)
                            small_step_2nd = small_step_tmp
                        else:
                            break
                if self.almost_connected(small_step_2nd, small_step_1st, threshold=step_length+2):
                    self._expand_node_and_edge(self._tree_1st, small_step_1st, small_step_2nd)
                    self.connected_node = small_step_2nd
                    print("succeed build two tree at epoch: {}".format(e))
                    return True
            if self._tree_1st.node_size > self._tree_2nd.node_size:
                self._tree_1st, self._tree_2nd = self._tree_2nd, self._tree_1st
        print("falied build trees in {} epoch!".format(max_epoch))
        return False

    def _solve_heuristic_to_goal(self, tree, goal):
        heurs = {}
        for node in tree.nodes():
            heurs[(node.nid, goal.nid)] = self.graph.distance(node, goal)
        return heurs

    def _search_path(self):
        if self.start == self._tree_1st.root:
            self.start_tree, self.goal_tree = self._tree_1st, self._tree_2nd
        else:
            self.start_tree, self.goal_tree = self._tree_2nd, self._tree_1st
        self.start_tree.heuristics = self._solve_heuristic_to_goal(self.start_tree, self.connected_node)
        self.goal_tree.heuristics = self._solve_heuristic_to_goal(self.goal_tree, self.connected_node)
        planner_1 = AStarPlanner()
        res_path_1 = planner_1.plan(self.start_tree, self.start, self.connected_node.nid)
        planner_2 = AStarPlanner()
        res_path_2 = planner_2.plan(self.goal_tree, self.goal, self.connected_node.nid)
        if res_path_1 is None or res_path_2 is None:
            print("184", type(res_path_1), type(res_path_2))
            print('187', self.start_tree.nodes())
            print('188', self.goal_tree.nodes())
            raise Exception("185 find none path")
            return self.result #is None
        self.result = res_path_1 + res_path_2.reverse()
        # print('174', type(res_path_1), type(res_path_2), type(res_path_2.reverse()), type(self.result))
        return self.result

    def plan(self, graph=None, start=None, goal=None, step_length=8, max_epoch=5000):
        if not self.check(graph, start, goal):
            return None
        if self._build_bi_tree(step_length, max_epoch):
            self._search_path()
        return self.result

if __name__ == "__main__":
    pass